/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.pinot.broker.routing.instanceselector;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.helix.model.ExternalView;
import org.apache.helix.model.IdealState;
import org.apache.pinot.broker.routing.segmentpreselector.SegmentPreSelector;
import org.apache.pinot.common.request.BrokerRequest;


/**
 * The instance selector selects server instances to serve the query based on the selected segments.
 */
public interface InstanceSelector {
  long NEW_SEGMENT_EXPIRATION_MILLIS = TimeUnit.MINUTES.toMillis(5);

  static boolean isNewSegment(long creationTimeMs, long currentTimeMs) {
    return isNewSegment(creationTimeMs, currentTimeMs, NEW_SEGMENT_EXPIRATION_MILLIS);
  }

  static boolean isNewSegment(long creationTimeMs, long currentTimeMs, long newSegmentExpirationMillis) {
    return creationTimeMs > 0 && currentTimeMs - creationTimeMs <= newSegmentExpirationMillis;
  }

  /**
   * Initializes the instance selector with the enabled instances, ideal state, external view and online segments
   * (segments with ONLINE/CONSUMING instances in the ideal state and pre-selected by the {@link SegmentPreSelector}).
   * Should be called only once before calling other methods.
   */
  void init(Set<String> enabledInstances, IdealState idealState, ExternalView externalView, Set<String> onlineSegments);

  /**
   * Processes the instances change. Changed instances are pre-computed based on the current and previous enabled
   * instances only once on the caller side and passed to all the instance selectors.
   */
  void onInstancesChange(Set<String> enabledInstances, List<String> changedInstances);

  /**
   * Processes the segment assignment (ideal state or external view) change based on the given online segments (segments
   * with ONLINE/CONSUMING instances in the ideal state and pre-selected by the {@link SegmentPreSelector}).
   */
  void onAssignmentChange(IdealState idealState, ExternalView externalView, Set<String> onlineSegments);

  /**
   * Selects the server instances for the given segments queried by the given broker request, returns a map from segment
   * to selected server instance hosting the segment and a set of unavailable segments (no enabled instance or all
   * enabled instances are in ERROR state).
   *
   * @param brokerRequest BrokerRequest for the query
   * @param segments segments for which instance needs to be selected
   * @param requestId requestId generated by the Broker for a query
   * @return instance of SelectionResult which describes the instance to pick for a given segment
   */
  SelectionResult select(BrokerRequest brokerRequest, List<String> segments, long requestId);

  /**
   * Returns the enabled server instances currently serving the table.
   */
  Set<String> getServingInstances();

  class SelectionResult {
    private final Pair<Map<String, String>, Map<String, String>/*optional segments*/> _segmentToInstanceMap;
    private final List<String> _unavailableSegments;
    private int _numPrunedSegments;

    public SelectionResult(Pair<Map<String, String>, Map<String, String>> segmentToInstanceMap,
        List<String> unavailableSegments, int numPrunedSegments) {
      _segmentToInstanceMap = segmentToInstanceMap;
      _unavailableSegments = unavailableSegments;
      _numPrunedSegments = numPrunedSegments;
    }

    /**
     * Returns the map from segment to selected server instance hosting the segment.
     */
    public Map<String, String> getSegmentToInstanceMap() {
      return _segmentToInstanceMap.getLeft();
    }

    /**
     * Returns the map from optional segment to selected server instance hosting the optional segment.
     * Optional segments can be skipped by broker or server upon any issue w/o failing the query.
     */
    public Map<String, String> getOptionalSegmentToInstanceMap() {
      return _segmentToInstanceMap.getRight();
    }

    /**
     * Returns the unavailable segments (no enabled instance or all enabled instances are in ERROR state).
     */
    public List<String> getUnavailableSegments() {
      return _unavailableSegments;
    }

    /**
     * Returns the number of segments pruned by the broker
     */
    public int getNumPrunedSegments() {
      return _numPrunedSegments;
    }

    /**
     * Sets the number of segments pruned by the broker
     */
    public void setNumPrunedSegments(int numPrunedSegments) {
      _numPrunedSegments = numPrunedSegments;
    }
  }
}
